# ---------BERT MASK----推理------------------
import pickle

import torch
from utils.datasets import BERTMlm
from utils.models import Ours
from torch.utils.data import DataLoader
from utils.functions import add_muti_mask_vec, drugseed_bert_ids, drug_names_bert_ids, BoxPlot
from transformers import BertTokenizer, BertForMaskedLM
from confs import step2 as conf
from tqdm import tqdm


def inf(model_name, device, topk, seedonly, type="inf"):
    if seedonly:
        dic = drugseed_bert_ids
    else:
        dic = drug_names_bert_ids
    conf.BATCHSIZE = 100
    res = []
    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
    model = Ours(conf).to(device)
    model.load_state_dict(torch.load(f"./models/step3train/{model_name}.pkl", map_location=device))
    infset = BERTMlm(conf, tokenizer, type=type)
    inf_loader = DataLoader(
        infset,
        batch_size=conf.BATCHSIZE,
        num_workers=4,
        pin_memory=True,
        shuffle=False
    )

    model.eval()
    with torch.no_grad():
        # for steps, _ in enumerate(tqdm(inf_loader)):
        for steps, _ in enumerate(inf_loader):
            for k in _[0].keys():
                _[0][k] = _[0][k].squeeze(1).to(device)

            output = model(_, type="inf")  # 这个写法在模型后加两层后就不够优雅了，但是为了统一代码结构还得这么写
            logits = output.logits
            mask_token_index = (_[0]["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]
            try:
                mask_logits_ = torch.stack([logits[i, j] for i, j in enumerate(mask_token_index)])
            except Exception as e:
                print(e)
                print(f"steps: {steps}")
            mask_logits = add_muti_mask_vec(mask_logits_, tokenizer)
            sorted, indices = torch.sort(mask_logits, descending=True)
            for i in range(indices.shape[0]):
                if len(set(dic.values()) & set(indices[i, :topk].tolist())) != 0:
                    res.append(infset.corp[steps * conf.BATCHSIZE + i])
    pickle.dump(res, open(f"./corpus/{step}/{model_name}_top{topk}_output.pkl", "wb"))


if __name__ == '__main__':
    device = "cuda:2"
    topk = 20
    seedonly = True
    step = "step3inf"
    model_name = "Epoch0_step4train_06_mlm"
    inf(model_name, device, topk, type="inf", seedonly=seedonly)
    res = pickle.load(open(f"./corpus/{step}/{model_name}_top{topk}_output.pkl", "rb"))
    print(f"{model_name}_top{topk}_output:{len(res)}")
